{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/multi_tenancy/multi_tenancy_rag.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-Tenancy RAG with LlamaIndex\n",
    "\n",
    "In this notebook you will look into building Multi-Tenancy RAG System using LlamaIndex.\n",
    "\n",
    "1. Setup\n",
    "2. Download Data\n",
    "3. Load Data\n",
    "4. Create Index\n",
    "5. Create Ingestion Pipeline\n",
    "6. Update Metadata and Insert documents\n",
    "7. Define Query Engines for each user\n",
    "8. Querying"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    " You should ensure you have `llama-index` and `pypdf` is installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index pypdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set OpenAI Key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"YOUR OPENAI API KEY\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import VectorStoreIndex\n",
    "from llama_index.vector_stores.types import MetadataFilters, ExactMatchFilter\n",
    "from llama_index import SimpleDirectoryReader\n",
    "from llama_index.ingestion import IngestionPipeline\n",
    "from llama_index.text_splitter import SentenceSplitter\n",
    "\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download Data\n",
    "\n",
    "We will use `An LLM Compiler for Parallel Function Calling` and `Dense X Retrieval: What Retrieval Granularity Should We Use?` papers for the demonstartions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-01-15 14:29:26--  https://arxiv.org/pdf/2312.04511.pdf\n",
      "Resolving arxiv.org (arxiv.org)... 151.101.131.42, 151.101.67.42, 151.101.3.42, ...\n",
      "Connecting to arxiv.org (arxiv.org)|151.101.131.42|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 755837 (738K) [application/pdf]\n",
      "Saving to: ‘llm_compiler.pdf’\n",
      "\n",
      "\r",
      "llm_compiler.pdf      0%[                    ]       0  --.-KB/s               \r",
      "llm_compiler.pdf    100%[===================>] 738.12K  --.-KB/s    in 0.004s  \n",
      "\n",
      "2024-01-15 14:29:26 (163 MB/s) - ‘llm_compiler.pdf’ saved [755837/755837]\n",
      "\n",
      "--2024-01-15 14:29:26--  https://arxiv.org/pdf/2312.06648.pdf\n",
      "Resolving arxiv.org (arxiv.org)... 151.101.131.42, 151.101.67.42, 151.101.3.42, ...\n",
      "Connecting to arxiv.org (arxiv.org)|151.101.131.42|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 1103758 (1.1M) [application/pdf]\n",
      "Saving to: ‘dense_x_retrieval.pdf’\n",
      "\n",
      "dense_x_retrieval.p 100%[===================>]   1.05M  --.-KB/s    in 0.005s  \n",
      "\n",
      "2024-01-15 14:29:26 (208 MB/s) - ‘dense_x_retrieval.pdf’ saved [1103758/1103758]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget --user-agent \"Mozilla\" \"https://arxiv.org/pdf/2312.04511.pdf\" -O \"llm_compiler.pdf\"\n",
    "!wget --user-agent \"Mozilla\" \"https://arxiv.org/pdf/2312.06648.pdf\" -O \"dense_x_retrieval.pdf\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = SimpleDirectoryReader(input_files=[\"dense_x_retrieval.pdf\"])\n",
    "documents_jerry = reader.load_data()\n",
    "\n",
    "reader = SimpleDirectoryReader(input_files=[\"llm_compiler.pdf\"])\n",
    "documents_ravi = reader.load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an Empty Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = VectorStoreIndex.from_documents(documents=[])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Ingestion Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = IngestionPipeline(\n",
    "    transformations=[\n",
    "        SentenceSplitter(chunk_size=512, chunk_overlap=20),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update Metadata and Insert Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for document in documents_jerry:\n",
    "    document.metadata[\"user\"] = \"Jerry\"\n",
    "\n",
    "nodes = pipeline.run(documents=documents_jerry)\n",
    "# Insert nodes into the index\n",
    "index.insert_nodes(nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for document in documents_ravi:\n",
    "    document.metadata[\"user\"] = \"Ravi\"\n",
    "\n",
    "nodes = pipeline.run(documents=documents_ravi)\n",
    "# Insert nodes into the index\n",
    "index.insert_nodes(nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Query Engines\n",
    "\n",
    "Define query engines for both the users with necessary filters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For Jerry\n",
    "jerry_query_engine = index.as_query_engine(\n",
    "    filters=MetadataFilters(\n",
    "        filters=[\n",
    "            ExactMatchFilter(\n",
    "                key=\"user\",\n",
    "                value=\"Jerry\",\n",
    "            )\n",
    "        ]\n",
    "    ),\n",
    "    similarity_top_k=3,\n",
    ")\n",
    "\n",
    "# For Ravi\n",
    "ravi_query_engine = index.as_query_engine(\n",
    "    filters=MetadataFilters(\n",
    "        filters=[\n",
    "            ExactMatchFilter(\n",
    "                key=\"user\",\n",
    "                value=\"Ravi\",\n",
    "            )\n",
    "        ]\n",
    "    ),\n",
    "    similarity_top_k=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">The paper mentions propositions as an alternative retrieval unit choice. Propositions are defined as atomic expressions of meanings in text that correspond to distinct pieces of meaning in the text. They are minimal and cannot be further split into separate propositions. Each proposition is contextualized and self-contained, including all the necessary context from the text to interpret its meaning. The paper demonstrates the concept of propositions using an example about the Leaning Tower of Pisa, where the passage is split into three propositions, each corresponding to a distinct factoid about the tower.</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Jerry has Dense X Rerieval paper and should be able to answer following question.\n",
    "response = jerry_query_engine.query(\n",
    "    \"what are propositions mentioned in the paper?\"\n",
    ")\n",
    "# Print response\n",
    "display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">LLMCompiler consists of three key components: an LLM Planner, a Task Fetching Unit, and an Executor. The LLM Planner identifies the execution flow by defining different function calls and their dependencies based on user inputs. The Task Fetching Unit dispatches the function calls that can be executed in parallel after substituting variables with the actual outputs of preceding tasks. Finally, the Executor executes the dispatched function calling tasks using the associated tools. These components work together to optimize the parallel function calling performance of LLMs.</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Ravi has LLMCompiler paper\n",
    "response = ravi_query_engine.query(\"what are steps involved in LLMCompiler?\")\n",
    "\n",
    "# Print response\n",
    "display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">The steps involved in LLMCompiler are not mentioned in the given context information.</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# This should not be answered as Jerry does not have information about LLMCompiler\n",
    "response = jerry_query_engine.query(\"what are steps involved in LLMCompiler?\")\n",
    "\n",
    "# Print response\n",
    "display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
